import pandas as pd
import io
from contextlib import redirect_stdout, redirect_stderr
import traceback
import ast
import streamlit as st
import numpy as np
import math
import random
import json
import re
import time
import copy
import plotly.express as px
import plotly.graph_objects as go

class CodeInterpreter:
    def __init__(self):
        self.allowed_modules = {
            'pandas': pd,
            'numpy': np,
            'math': math,
            'random': random,
            'json': json,
            're': re,
            'time': time,
            'copy': copy,
            'px': px,
            'go': go
        }
        
    def is_safe_code(self, code: str) -> bool:
        """检查代码是否安全可执行"""
        try:
            tree = ast.parse(code)
            for node in ast.walk(tree):
                # 检查是否有危险的内置函数调用
                if isinstance(node, ast.Call):
                    if isinstance(node.func, ast.Name):
                        if node.func.id in ['eval', 'exec', 'compile', '__import__']:
                            return False
                # 检查导入语句
                elif isinstance(node, (ast.Import, ast.ImportFrom)):
                    for name in node.names:
                        module = name.name.split('.')[0]
                        if module not in self.allowed_modules:
                            return False
            return True
        except SyntaxError:
            return False

    def prepare_environment(self, df=None):
        """准备代码执行环境"""
        env = {
            'pd': pd,
            'df': df,
            'st': st,
            'np': np,
            'math': math,
            'random': random,
            'json': json,
            're': re,
            'time': time,
            'copy': copy,
            'px': px,
            'go': go
        }
        return env

    def execute_code(self, code: str, df=None):
        """执行代码并返回结果"""
        if not self.is_safe_code(code):
            return {
                'success': False,
                'error': "代码包含不允许的操作或模块",
                'output': None,
                'df': df
            }

        output = io.StringIO()
        local_dict = self.prepare_environment(df)

        try:
            with redirect_stdout(output), redirect_stderr(output):
                exec(code, local_dict)

            return {
                'success': True,
                'error': None,
                'output': output.getvalue(),
                'df': local_dict.get('df', df),
                'fig': local_dict.get('fig', None)
            }
        except Exception as e:
            return {
                'success': False,
                'error': traceback.format_exc(),
                'output': output.getvalue(),
                'df': df
            }

    def format_code(self, code: str) -> str:
        """格式化代码，移除多余的空行和注释"""
        lines = code.split('\n')
        formatted_lines = []
        for line in lines:
            stripped = line.strip()
            if stripped and not stripped.startswith('#'):
                formatted_lines.append(line)
        return '\n'.join(formatted_lines) 